//! Page table management in Zig
//! SPDX-License-Identifier: MPL-2.0

const std = @import("std");
const lib = @import("lib.zig");
const frame = @import("frame.zig");

const Vaddr = lib.Vaddr;
const Paddr = lib.Paddr;
const PageFlags = lib.PageFlags;
const PagingLevel = lib.PagingLevel;
const PAGE_SIZE = lib.PAGE_SIZE;
const MemoryError = lib.MemoryError;

// Page table entry
pub const PageTableEntry = packed struct {
    present: bool = false,
    writable: bool = false,
    user_accessible: bool = false,
    write_through: bool = false,
    cache_disabled: bool = false,
    accessed: bool = false,
    dirty: bool = false,
    huge_page: bool = false,
    global: bool = false,
    available1: u3 = 0,
    address: u40 = 0, // Physical address bits [51:12]
    available2: u11 = 0,
    no_execute: bool = false,

    const Self = @This();

    pub fn init(phys_addr: Paddr, flags: PageFlags) Self {
        return Self{
            .present = flags.present,
            .writable = flags.writable,
            .user_accessible = flags.user_accessible,
            .write_through = flags.write_through,
            .cache_disabled = flags.cache_disabled,
            .accessed = flags.accessed,
            .dirty = flags.dirty,
            .huge_page = flags.huge_page,
            .global = flags.global,
            .no_execute = flags.no_execute,
            .address = @truncate(phys_addr >> 12),
        };
    }

    pub fn paddr(self: Self) Paddr {
        return (@as(Paddr, self.address) << 12);
    }

    pub fn isPresent(self: Self) bool {
        return self.present;
    }

    pub fn isHugePage(self: Self) bool {
        return self.huge_page;
    }

    pub fn clear(self: *Self) void {
        self.* = Self{};
    }
};

// Page table structure
pub const PageTable = struct {
    entries: [512]PageTableEntry,
    frame_ref: ?frame.Frame,

    const Self = @This();

    pub fn init() Self {
        return Self{
            .entries = [_]PageTableEntry{PageTableEntry{}} ** 512,
            .frame_ref = null,
        };
    }

    pub fn initWithFrame(table_frame: frame.Frame) Self {
        const vaddr = lib.paddrToVaddr(table_frame.startPaddr());
        const entries_ptr: *[512]PageTableEntry = @ptrFromInt(vaddr);
        
        return Self{
            .entries = entries_ptr.*,
            .frame_ref = table_frame,
        };
    }

    pub fn getEntry(self: *Self, index: usize) ?*PageTableEntry {
        if (index >= 512) return null;
        return &self.entries[index];
    }

    pub fn setEntry(self: *Self, index: usize, paddr: Paddr, flags: PageFlags) !void {
        if (index >= 512) return MemoryError.InvalidAddress;
        self.entries[index] = PageTableEntry.init(paddr, flags);
    }

    pub fn clearEntry(self: *Self, index: usize) !void {
        if (index >= 512) return MemoryError.InvalidAddress;
        self.entries[index].clear();
    }

    pub fn vaddrToIndex(vaddr: Vaddr, level: PagingLevel) usize {
        const shift = 12 + (level - 1) * 9;
        return (vaddr >> shift) & 0x1ff;
    }
};

// Page table walker for address translation
pub const PageTableWalker = struct {
    root_table: *PageTable,

    const Self = @This();

    pub fn init(root_table: *PageTable) Self {
        return Self{ .root_table = root_table };
    }

    pub fn translate(self: Self, vaddr: Vaddr) ?Paddr {
        var current_table = self.root_table;
        var level: PagingLevel = lib.NR_LEVELS;

        while (level > 1) {
            const index = PageTable.vaddrToIndex(vaddr, level);
            const entry = current_table.getEntry(index) orelse return null;

            if (!entry.isPresent()) return null;
            
            if (entry.isHugePage()) {
                // Handle huge page
                const page_size = lib.pageSize(level);
                const offset = vaddr & (page_size - 1);
                return entry.paddr() + offset;
            }

            // Navigate to next level
            const next_table_paddr = entry.paddr();
            const next_table_vaddr = lib.paddrToVaddr(next_table_paddr);
            current_table = @ptrFromInt(next_table_vaddr);
            level -= 1;
        }

        // Final level translation
        const index = PageTable.vaddrToIndex(vaddr, 1);
        const entry = current_table.getEntry(index) orelse return null;
        
        if (!entry.isPresent()) return null;
        
        const offset = vaddr & (PAGE_SIZE - 1);
        return entry.paddr() + offset;
    }

    pub fn map(self: Self, vaddr: Vaddr, paddr: Paddr, flags: PageFlags, frame_allocator: frame.FrameAllocator, allocator: std.mem.Allocator) !void {
        var current_table = self.root_table;
        var level: PagingLevel = lib.NR_LEVELS;

        while (level > 1) {
            const index = PageTable.vaddrToIndex(vaddr, level);
            const entry = current_table.getEntry(index) orelse return MemoryError.InvalidAddress;

            if (!entry.isPresent()) {
                // Allocate new page table
                const options = frame.FrameAllocOptions.new().withZeroed(true);
                const new_frame = try frame_allocator.allocFrame(options, allocator);
                const new_table_paddr = new_frame.startPaddr();
                
                const table_flags = PageFlags{
                    .present = true,
                    .writable = true,
                    .user_accessible = flags.user_accessible,
                };
                
                entry.* = PageTableEntry.init(new_table_paddr, table_flags);
            }

            if (level == 1 or !entry.isHugePage()) {
                const next_table_paddr = entry.paddr();
                const next_table_vaddr = lib.paddrToVaddr(next_table_paddr);
                current_table = @ptrFromInt(next_table_vaddr);
            }
            
            level -= 1;
        }

        // Map at final level
        const index = PageTable.vaddrToIndex(vaddr, 1);
        try current_table.setEntry(index, paddr, flags);
    }

    pub fn unmap(self: Self, vaddr: Vaddr) !void {
        var current_table = self.root_table;
        var level: PagingLevel = lib.NR_LEVELS;

        while (level > 1) {
            const index = PageTable.vaddrToIndex(vaddr, level);
            const entry = current_table.getEntry(index) orelse return MemoryError.InvalidAddress;

            if (!entry.isPresent()) return MemoryError.NotMapped;
            
            if (entry.isHugePage()) {
                entry.clear();
                return;
            }

            const next_table_paddr = entry.paddr();
            const next_table_vaddr = lib.paddrToVaddr(next_table_paddr);
            current_table = @ptrFromInt(next_table_vaddr);
            level -= 1;
        }

        // Unmap at final level
        const index = PageTable.vaddrToIndex(vaddr, 1);
        try current_table.clearEntry(index);
    }
};

// Tests
test "page table entry" {
    const testing = std.testing;
    const flags = PageFlags{
        .present = true,
        .writable = true,
        .user_accessible = false,
    };
    
    const entry = PageTableEntry.init(0x1000, flags);
    try testing.expect(entry.isPresent());
    try testing.expectEqual(@as(Paddr, 0x1000), entry.paddr());
}

test "vaddr to index calculation" {
    const testing = std.testing;
    const vaddr: Vaddr = 0x12345000;
    
    // Test different levels
    const index_l1 = PageTable.vaddrToIndex(vaddr, 1);
    const index_l2 = PageTable.vaddrToIndex(vaddr, 2);
    
    try testing.expect(index_l1 < 512);
    try testing.expect(index_l2 < 512);
}